﻿import pygame, os, time, thread
from pygame.locals import *
from Tkinter import *
from tkFileDialog   import asksaveasfilename

# import modules such as task, queue, cpu from modules folder
path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../modules'))
if not path in sys.path:
  sys.path.insert(1, path)
del path

from constants import *
from task import Task
from queue import Queue
from vlist import List
from cpu import CPU
from simulator import *

#####################################################
#                                                   #
#  NewSimulation class:                             # 
#                                                   # 
#  This class allows the user to create             #
#  custom simulations using tasks, queues and lists #
#                                                   # 
#  Author: Felix Rohrer                             #     
#  Course: CS MET 575                               #
#  Year:   2011                                     #
#                                                   #
#####################################################

class NewSimulation():
  def __init__(self, mainMenu=None, tasks=[]):

    # main menu object
    self.mainMenu = mainMenu

    # helper menu (allows the user to choose data structure names and lengths)
    self.helper = None

    self.frame = mainMenu.ctrlRegion

    # draw area color
    self.drawAreaColor = DRAWAREA_BG_COLOR

    # label to display informations
    self.infoMsg = StringVar()
    self.infoMsgLabel = Label(self.frame, textvariable=self.infoMsg, padx=10, justify=LEFT, anchor=W, relief = GROOVE, width=32)
    
    # number of elements used
    self.numTasks = 0
    self.numQueues = 0
    self.numLists = 0
    self.numCPUs = 0

    # this variable indicates, if simulation in process
    self.going = False
  
    # recording related variables
    self.recording = False
    self.flow = ''
    self.flow_len = len(self.flow)
    self.nop_time = 0
    self.fileName = StringVar()

    # all tasks
    self.tasks = []
    self.ds = {}

    # task names
    self.taskNameFields = []

    # queue related variables
    self.queueNameFields = []
    self.queueLengths = []

    # list related variables
    self.listNameFields = []
    self.listLengths = []

    self.cpuNameFields = []

  # display information string in main window
  def displayInfo(self,infoMsg, pos=(0,20,3),color=ERR_BG_COLOR):
    if infoMsg is None:
      self.infoMsg.set('')
      self.infoMsgLabel.config(bg=MAINMENU_BG_COLOR,relief=FLAT)
    else:
      self.infoMsg.set(infoMsg)
      self.infoMsgLabel.config(bg=color,relief=GROOVE)
    self.infoMsgLabel.grid(column=pos[0],row=pos[1],columnspan=pos[2],sticky=W)

  # every simulation needs to provide a controls function
  # the controls function will be called from the main menu, 
  # once this specific simulation is requested
  def controls(self):
    self.frame.grid()
    self.frame.grid_columnconfigure(0,minsize=130)

    self.selectDSLabel = Label(self.frame, text="Choose objects: ", height=2,font=TITLE_FONT)
    self.controlLabel = Label(self.frame, text="Control simulation: ", height=2,font=TITLE_FONT)
    
    # information label
    self.infoLabel = Label(self.frame, text="Information:", height=2,font=TITLE_FONT)

    self.startButton = Button(self.frame, text="New simulation", command=self.createObjects,width=15)
    self.closeButton = Button(self.frame, text="Close simulation", command=self.closeSimulation, width=15,state=DISABLED)

    # show queue slider
    self.queueLabel = Label(self.frame, text="Queues ")
    self.queueScale = Scale(self.frame, from_=0,to=5, orient=HORIZONTAL, command=self.setNumQueues)

    # show list slider
    self.listLabel = Label(self.frame, text="Lists ")
    self.listScale = Scale(self.frame, from_=0,to=5, orient=HORIZONTAL, command=self.setNumLists)

    # show list slider
    self.cpuLabel = Label(self.frame, text="CPUs ")
    self.cpuScale = Scale(self.frame, from_=0,to=5, orient=HORIZONTAL, command=self.setNumCPUs)

    # show task slider
    self.taskLabel = Label(self.frame, text="Tasks ")
    self.taskScale = Scale(self.frame, from_=1,to=10, orient=HORIZONTAL, command=self.setNumTasks)

    # record button
    self.startRecordButton = Button(self.frame, text="Start recording", command=self.startRecording, foreground = 'red', activeforeground = 'red', state=DISABLED,width=15)
    self.stopRecordButton = Button(self.frame, text="Stop recording", command=self.stopRecording, state=DISABLED,width=15)

    # file save
    self.nameLabel = Label(self.frame, text="Save simulation: ", height=2,font=TITLE_FONT)
    self.nameField = Entry(self.frame, width=37,textvariable=self.fileName,state=DISABLED)
    self.chooseButton = Button(self.frame, text="Choose", command=self.saveFile,width=15)
    self.saveSimButton = Button(self.frame, text="Save", command=self.saveRecording, state=DISABLED,width=15)

    # show selection label
    self.selectDSLabel.grid(columnspan=2,column=0,row=0,sticky=W)

    # queues
    self.queueLabel.grid(column=0, row=2, sticky=W)
    self.queueScale.grid(column=1, row=2, sticky=W)

    # list
    self.listLabel.grid(column=0, row=3, sticky=W)
    self.listScale.grid(column=1, row=3, sticky=W)

    # cpu
    self.cpuLabel.grid(column=0, row=4, sticky=W)
    self.cpuScale.grid(column=1, row=4, sticky=W)

    # tasks
    self.taskLabel.grid(column=0, row=5, sticky=W)
    self.taskScale.grid(column=1, row=5, sticky=W)

    # start simulation
    self.controlLabel.grid(columnspan=2,column=0,row=6,sticky=W)
    self.startButton.grid(column=0,row=7,sticky=W)
    self.closeButton.grid(column=1,row=7,sticky=W)

    self.startRecordButton.grid(column=0, row=8, sticky=W)
    self.stopRecordButton.grid(column=1, row=8, sticky=W)

    # save as file
    self.nameLabel.grid(column=0, row=9, sticky=W)
    self.nameField.grid(columnspan=2,column=0, row=10, sticky=W)
    self.chooseButton.grid(column=0, row=11, sticky=W)
    self.saveSimButton.grid(column=1, row=11, sticky=E)

    # show information
    self.infoLabel.grid(columnspan=2,column=0,row=12,sticky=W)
    self.displayInfo("Choose your objects and then \nclick 'Start Simulation'. ",color=INFO_BG_COLOR)

  # destroy the 'new simulation' control panel with all 
  # its dependent windows
  def destroy(self):

    # close the simulation window if opened
    self.closeSimulation()

    if self.helper is not None:
      self.helper.destroy()
      self.helper = None    

    # wait until simulation stopped
    while self.mainMenu.simulationRunning:
      pass

    # do not show any information
    self.displayInfo(None)

    self.frame = None
    self.mainMenu = None

  # set amount of numbers
  def setNumTasks(self,upd):
    self.numTasks = int(upd)

  # set amount of queues used
  def setNumQueues(self,upd):
    self.numQueues = int(upd)

  # set amount of queues used
  def setNumLists(self,upd):
    self.numLists = int(upd)

  # set amount of queues used
  def setNumCPUs(self,upd):
    self.numCPUs = int(upd)

  # open a file dialog that asks where to store the current simulation
  def saveFile(self):
    f = ''
    try:
      f = asksaveasfilename(initialdir=SIMFOLDER,filetypes=[('sim Files','*.sim')])
    except:
      self.displayInfo("Error: could not save simulation. ")
      return
    if f is not '':
      self.fileName.set(f)
      self.saveSimButton.config(state=NORMAL)
    else:
      self.displayInfo("Error: file is empty/malformed. ")

  # start recording
  def startRecording(self):
    if self.mainMenu.simulationRunning:
      # set initial data
      self.flow = ''

      # add all tasks with their position
      for t in range(len(self.tasks)):
        self.flow += 't'+str(t)+self.tasks[t].name+'-0-'+str(self.tasks[t].pos)+';'

      # add all data structures with their position 
      for k,v in self.ds.iteritems():
        if k[0] == 'q':
          self.flow += k+v.name+'-'+str(v.getLength())+'-'+str(v.pos)+';'
        elif k[0] == 'l':
          self.flow += k+v.name+'-'+str(v.getLength())+'-'+str(v.pos)+';'
        elif k[0] == 'c':
          self.flow += k+v.name+'-'+str(v.getLength())+'-'+str(v.pos)+';'
  
      self.flow_len = len(self.flow)

      # set mark
      self.flow += 'END_INITIAL'

      self.recording = True
      self.displayInfo("Recording started. ",color=INFO_BG_COLOR)
      self.stopRecordButton.config(state=NORMAL)
    else:
      self.displayInfo("Error: no simulation running yet, \nstart simulation first. ")

  # save recording
  def saveRecording(self):
    self.saveSimButton.config(state=NORMAL)
    fileName = self.nameField.get()
    if self.flow is '':
      self.displayInfo("Error: no simulation recorded. ")
      return
    if fileName is '':
      self.displayInfo("Error: no filename selected. ")
      self.nameField.config(background=ERR_BG_COLOR)
      return
    self.displayInfo("Success: simulation saved. ", color=INFO_BG_COLOR)
    fh = open(fileName, 'w')
    fh.write(self.flow)
    fh.close()

  # stop recording
  def stopRecording(self):
    if self.mainMenu.simulationRunning:
      if self.recording:
        self.recording = False 
        self.displayInfo("Stopped recording. ",color=INFO_BG_COLOR)
      else:
        self.displayInfo("Error: not recording. ")
    else:
      self.displayInfo("Error: no simulation running yet, \nstart simulation first. ")

  # this function creates the involved items and names them 
  # according to the requested names
  def createObjects(self):
    self.helper = Toplevel()
    self.helper.title("Customize Simulation")
    self.helper.grid()

    self.objNameLabel = Label(self.helper, text="Define names (or leave empty)", height=2,font=HELPER_FONT)
    self.objNameLabel.grid(columnspan=2,column=0,row=0,sticky=W)

    # get correct amount of tasks, queues and lists
    self.numTasks = self.taskScale.get() 
    self.numQueues = self.queueScale.get() 
    self.numLists = self.listScale.get() 
    self.numCPUs = self.cpuScale.get() 

    # remove all previously used name and length fields
    del self.taskNameFields[:]
    del self.queueNameFields[:]
    del self.listNameFields[:]
    del self.cpuNameFields[:]
    del self.queueLengths[:]
    del self.listLengths[:]

    # create required amount of tasks
    for p in range(self.numTasks):
      label = Label(self.helper, text="Task "+str(p)+": ")
      field = Entry(self.helper, width=2)

      # insert default value
      field.insert(0,"T"+str(p))

      # textfields for tasks can be split into 2 columns to save space
      # max tasks currently is 10, so 
      label.grid(column=(p/5)*2, row=p%5+1, sticky=W)
      field.grid(column=(p/5)*2+1, row=p%5+1, sticky=W)
      self.taskNameFields.append(field)

    # create entry fields for queue names
    for q in range(self.numQueues):
      label = Label(self.helper, text="Queue "+str(q)+": ")
      field = Entry(self.helper, width=10)

      # insert default value
      field.insert(0,"Queue "+str(q))

      lengthLabel = Label(self.helper, text="Length: ")
      queueLength = Scale(self.helper, from_=1,to=10, orient=HORIZONTAL)
      queueLength.set(5)

      label.grid(column=0, row=self.numTasks+q+1, sticky=W)
      field.grid(column=1, row=self.numTasks+q+1, sticky=W)
      lengthLabel.grid(column=2, row=self.numTasks+q+1, sticky=W)
      queueLength.grid(column=3, row=self.numTasks+q+1, sticky=W)
      self.queueNameFields.append(field)
      self.queueLengths.append(queueLength)

    # create entry fields for list names
    for l in range(self.numLists):
      label = Label(self.helper, text="List "+str(l)+": ")
      field = Entry(self.helper, width=10)

      # insert default value
      field.insert(0,"List "+str(l))

      lengthLabel = Label(self.helper, text="Length: ")
      listLength = Scale(self.helper, from_=1,to=10, orient=HORIZONTAL)
      listLength.set(5)

      label.grid(column=0, row=self.numTasks+self.numQueues+l+1, sticky=W)
      field.grid(column=1, row=self.numTasks+self.numQueues+l+1, sticky=W)
      lengthLabel.grid(column=2, row=self.numTasks+self.numQueues+l+1, sticky=W)
      listLength.grid(column=3, row=self.numTasks+self.numQueues+l+1, sticky=W)
      self.listNameFields.append(field)
      self.listLengths.append(listLength)

    # create entry fields for CPU names
    for c in range(self.numCPUs):
      label = Label(self.helper, text="CPU "+str(c)+": ")
      field = Entry(self.helper, width=10)

      # insert default value
      field.insert(0,"CPU "+str(c))

      label.grid(column=0, row=self.numTasks+self.numQueues+self.numLists+c+1, sticky=W)
      field.grid(column=1, row=self.numTasks+self.numQueues+self.numLists+c+1, sticky=W)
      self.cpuNameFields.append(field)

    startSim = Button(self.helper, text="Start", command=self.startSimulation)
    startSim.grid(column=0, row=self.numTasks+self.numQueues+self.numLists+self.numCPUs+2, sticky=W)

  # close simulation window
  def closeSimulation(self):
    self.going = False

  # every simulation needs this function in order to create
  # a new thread that will start the actual simulation window
  # A new thread is required to keep the main menu responsive
  def startSimulation(self):
    # only start new simulation if no other simulation running
    if not self.mainMenu.simulationRunning:

      # empty tasks list
      tasks = []

      # create required amount of tasks
      for p in range(self.numTasks):
        name = self.taskNameFields[p].get()
        if name is '':
          name = str(p)
        tasks.append(Task(name,pos=(5+p*25,150)))

      # empty data structure dictionary
      ds = {}
    
      # create required amount of queues
      for q in range(self.numQueues):
        name = self.queueNameFields[q].get()
        value = self.queueLengths[q].get()
        if name is '':
          name = str(q)
        ds['q'+str(q)] = Queue(value, [25+q*10,(q+1)*10], name)

      # create required amount of queues
      for l in range(self.numLists):
        name = self.listNameFields[l].get()
        value = self.listLengths[l].get()
        if name is '':
          name = str(l)
        ds['l'+str(l)] = List(value, [25+l*10,200+(l+1)*10], name)

      # create required amount of queues
      for c in range(self.numCPUs):
        name = self.cpuNameFields[c].get()
        if name is '':
          name = str(c)
        ds['c'+str(c)] = CPU(name,[25+c*10,350+(c+1)*10])

      # helper menu not needed anymore, destroy it
      self.helper.destroy()
      self.helper = None

      # start new thread to run simulation
      # -> keeps main menu responsive
      thread.start_new_thread(self.plot,(tasks,ds,))
      self.mainMenu.simulationRunning = True
      self.startRecordButton.config(state=NORMAL)
      self.displayInfo("1. Arrange elements\n2. Press 'Start recording' to record. \n3. Drag and drop tasks.\n4. When finished, click 'Stop recording'.\n5. Save recorded simulation. ",color=INFO_BG_COLOR)
    else:
      self.displayInfo("Error, another simulation already running. ")

  # update recording flow
  def updateFlow(self,new=None,newType='PAUSE'):
    # p = pause
    if self.nop_time > 1:
      self.flow += 'p'+str(self.nop_time)+';'
    if not new == None:
      if newType == 'COORD':
        # c = coordinate
        self.flow += 'c'
      elif newType == 'ADDTASK':
        # a = Add task
        self.flow += 'a'
      elif newType == 'DROPTASK':
        # d = Drop task
        self.flow += 'd'
      self.flow += str(new)+';'
    self.flow_len = len(self.flow)
    self.nop_time = 0

  # returns False if task could not be removed
  def removeFromCurrentQueue(self, task):
    if task.getDSType() == 'QUEUE':
      if task.currentDS.peakNext() == task:
        task.currentDS.pop()
      else:
        # in a queue, only first element can be removed
        return False
    elif task.getDSType() == 'LIST':
      task.currentDS.remove(task)
    return True

  # check collision with any of the data structures
  def collisionWithDS(self, draggedTask):
    currRect = Rect(draggedTask.pos[0],draggedTask.pos[1],draggedTask.size[0],draggedTask.size[1])
    found = False
    for v in self.ds.itervalues():
      listPos = currRect.collidelist(v.rects)
      # currently dragged task is colliding with data structure
      if listPos >= 0:
        # get data structure type (queue or list)
        dst = v.DSType

        # if currently dragged task can be removed from its
        # old data structure (or does not belong to a data structure)
        # add the currently dragged task to this new data structure
        if self.removeFromCurrentQueue(draggedTask):
	  if dst == 'QUEUE':
	    if v.push(draggedTask):
	      found = True
            else:
              self.displayInfo("ERROR: Could not add task to "+v.name)
	  elif dst == 'LIST':
	    v.insert(draggedTask,listPos)
	    found = True
    # if currently dragged task did not collide with 
    # any data structure, the task is on a non-data structure
    # field and its current data structure should be removed
    if not found:
      self.removeFromCurrentQueue(draggedTask)

  # plot the graph and all additional buttons and input fields 
  def plot(self,tasks=[],ds={}):

    # create clock to limit FPS
    clock = pygame.time.Clock()

    os.environ['SDL_VIDEO_WINDOW_POS'] = DRAWAREA_INIT_POS

    self.tasks = tasks
    self.ds = ds

    # initialize simulation window
    pygame.init()
    pygame.display.init()

    self.going = True
    self.screenDim = {'width': DRAWAREA_WIDTH, 'height': DRAWAREA_HEIGHT}

    self.screen = pygame.display.set_mode ((self.screenDim['width'], self.screenDim['height']))
    pygame.display.set_caption("Custom simulation")

    background = pygame.Surface(self.screen.get_size())
    background = background.convert()

    # drawing area
    drawSize = self.screenDim['width'],self.screenDim['height']
    drawArea = pygame.Surface(drawSize)
    drawArea.fill(self.drawAreaColor)

    # define font
    textFont = pygame.font.SysFont(SYS_FONT,SYS_FONT_SIZE)
    textFont.set_bold(True)

    # signal a mouse click and hold event
    mouseDown = False
    draggedTask = None
    draggedDS = None

    # enable and disable buttons
    self.closeButton.config(state=NORMAL)
    self.startButton.config(state=DISABLED)

    # sort data structures from dictionary
    sortedDS = self.ds.keys()
    sortedDS.sort()

    while self.going:

      # update nop time
      if self.recording:
	if len(self.flow) == self.flow_len:
	  self.nop_time +=1

      # limit frame rate
      time_passed = clock.tick(FPS)

      event = pygame.event.poll()
      if event.type == QUIT:
        self.going = False
      if event.type == KEYDOWN:
        if event.key == K_ESCAPE:
          self.going = False
      if event.type == MOUSEBUTTONDOWN:
        # check if left mouse pressed down
        if pygame.mouse.get_pressed()[0] == True:
          # left mouse clicked and held
          mouseDown = True
      # mouse released
      if event.type == MOUSEBUTTONUP:
        # mouse click released
        mouseDown = False
       
        # stop moving any previously moved task
        if draggedTask is not None:

          taskCond = draggedTask.getCondition()

          # if current task is not allowed to move
          # an error message was written when this task
          # was picked up. 
          # The task is dropped here, so the error message
          # has to be removed
          if taskCond == 'NOMOVE':
            self.displayInfo(None)

          # NOMOVE tasks are ignored by the recording
          if self.recording and draggedTask.getCondition() is not 'NOMOVE':
            self.updateFlow(draggedTask.name,'DROPTASK')

          draggedTask.dropped()

          # check for collision with queues and lists
          self.collisionWithDS(draggedTask)

          draggedTask = None

        if not draggedDS == None:
          draggedDS.dropped()
          draggedDS = None

      # draw graph
      self.screen.blit(background, (0, 0))

      # draw drawing area
      self.screen.blit(drawArea, (0,0))

      # update the drawing area
      drawArea.fill(self.drawAreaColor)
      pygame.draw.rect(drawArea, (0,0,0), (0,0,drawSize[0],drawSize[1]),1)

      # update currently dragged tasks position
      if mouseDown:
        newPos = pygame.mouse.get_pos()

        # check if mouse click collides with any task or 
        # any task is currently already selected to be moved
        if draggedTask is None:
          # only look for new task to drag, if no data structure
          # is currently dragged
          if draggedDS is None:
	    for t in self.tasks:
	      if t.collision(newPos):
		draggedTask = t
		# if the current task is in a queue, only allow it to be
		# moved if the task is at the queues head
		if draggedTask.getDSType() == 'QUEUE':
		  if draggedTask == draggedTask.currentDS.peakNext():
                    # only add task in queue to recording
                    # if task is at the head of the queue
		    if self.recording:
		      self.updateFlow(str(t.name),'ADDTASK')
		    draggedTask.setCondition('DRAGGED')
		  else:
		    draggedTask.setCondition('NOMOVE')
                # add this task (which is not in queue) to recording
                elif self.recording:
		  self.updateFlow(str(t.name),'ADDTASK')
		break
          # if draggedtask is still none, check collision
          # with a queue, list or CPU
          # only allow movement of data structures while
          # not recording
          if not self.recording:
	    if draggedTask is None and draggedDS is None:
	      for q in self.ds.itervalues():
		if q.collision(newPos):
		  draggedDS = q
        else:
          oldPos = draggedTask.getPos(offset=True)
          # update tasks position if different from previous position
          if not newPos[0] == oldPos[0] or not newPos[1] == oldPos[1]:
            taskCond = draggedTask.getCondition()
            if self.recording and taskCond is not 'NOMOVE':
              self.updateFlow(str(newPos),'COORD')
            if taskCond is 'NOMOVE':
              self.displayInfo("ERROR: task not head of queue,\n"+" "*12+"can't remove. ")
	    draggedTask.setPos(newPos,dragged=True)
        
        # move data structure if one is currently selected
        if not draggedDS == None:
          draggedDS.setPos(newPos)

      # draw all data structures
      for k in sortedDS:
        self.ds[k].display(drawArea)

      # draw all tasks
      for t in self.tasks:
        if not t == draggedTask:
          t.display(drawArea)

      # this ensures, that the currently dragged
      # task will always be displayed last
      # -> will be drawn on top of other objects
      if not draggedTask == None:
        draggedTask.display(drawArea)

      pygame.display.flip()

    # add last nop time to flow
    if self.nop_time > 1:
      self.flow += 'p'+str(self.nop_time)

    self.displayInfo("Stopping simulation...")
    self.frame.update_idletasks()

    # restore task counter to 0
    Task.TASK_COUNT = 0

    # uninitialize queue object
    del self.tasks[:]
    del self.taskNameFields[:]
    self.ds.clear()
 
    # enable and disable correct buttons
    self.closeButton.config(state=DISABLED)
    self.startButton.config(state=NORMAL)
    self.stopRecordButton.config(state=DISABLED)
    self.startRecordButton.config(state=DISABLED)

    pygame.quit ()
    self.displayInfo("Choose your objects and then \nclick 'Start Simulation'. ",color=INFO_BG_COLOR)
    self.mainMenu.simulationRunning = False
